from llmConfig import get_embeddings
from pymilvus import Collection


def acupointTool(query: str):
    try:
        # 加载集合
        collection = Collection("acupoint")
        collection.load()

        # 获取查询的向量表示
        query_vector = get_embeddings(query)

        # 执行向量搜索，返回最相关的结果
        search_results = collection.search(
            data=[query_vector],
            anns_field="vector",
            param={"metric_type": "L2", "params": {"nprobe": 10}},  # L2距离度量，nprobe控制搜索范围
            limit=500,  # 限制返回1个最相关的结果
            output_fields=["metadata", "text"]  # 返回的字段：元数据和文本
        )

        # 格式化输出，返回相关的内容
        formatted_context = ""
        for result in search_results[0]:  # 搜索结果中可能有多个匹配项
            metadata = result.entity.get('metadata')
            text = result.entity.get('text')

            # 将相关元数据和文本整合成返回内容
            formatted_context += f"穴位名称: {metadata.get('name', 'N/A')}\n"
            formatted_context += f"穴位位置: {metadata.get('location', 'N/A')}\n"
            formatted_context += f"操作方法: {metadata.get('operation', 'N/A')}\n"
            formatted_context += f"描述: {text}\n"

        # 如果没有匹配到结果，返回相应提示
        if not formatted_context:
            return "Sorry, no relevant acupoint found."

        return formatted_context

    except Exception as e:
        print(f"Error during the search operation: {e}")
        return "Sorry, an error occurred while processing your request."

#获取穴位
def acu_nameTool(query: str):
    try:
        collection = Collection("acupoint_name")
        collection.load()

        query_vector = get_embeddings(query)

        search_results = collection.search(
            data=[query_vector],
            anns_field="vector",
            param={"metric_type": "L2", "params": {"nprobe": 10}},
            limit=500,
            output_fields=["metadata", "text"]
        )

        formatted_context = "\n\n".join(
            [f"{doc.entity.get('text')}" for doc in search_results[0]]
        )

        return str(formatted_context)

    except Exception as e:
        print(f"Error during the search operation: {e}")
        return "Sorry to find nothing"

#获取身体部位
def acu_locTool(query: str):
    try:
        collection = Collection("acupoint_location")
        collection.load()

        query_vector = get_embeddings(query)

        search_results = collection.search(
            data=[query_vector],
            anns_field="vector",
            param={"metric_type": "L2", "params": {"nprobe": 10}},
            limit=500,
            output_fields=["metadata", "text"]
        )

        formatted_context = "\n\n".join(
            [f"{doc.entity.get('text')}" for doc in search_results[0]]
        )

        return str(formatted_context)

    except Exception as e:
        print(f"Error during the search operation: {e}")
        return "Sorry to find nothing"

#获取疾病
def diseTool(query: str):
    try:
        collection = Collection("disease")
        collection.load()

        query_vector = get_embeddings(query)

        search_results = collection.search(
            data=[query_vector],
            anns_field="vector",
            param={"metric_type": "L2", "params": {"nprobe": 10}},
            limit=500,
            output_fields=["metadata", "text"]
        )

        formatted_context = "\n\n".join(
            [f"{doc.entity.get('text')}" for doc in search_results[0]]
        )

        return str(formatted_context)

    except Exception as e:
        print(f"Error during the search operation: {e}")
        return "Sorry to find nothing"

#获取症状
def sympTool(query: str):
    try:
        collection = Collection("symptom")
        collection.load()

        query_vector = get_embeddings(query)

        search_results = collection.search(
            data=[query_vector],
            anns_field="vector",
            param={"metric_type": "L2", "params": {"nprobe": 10}},
            limit=500,
            output_fields=["metadata", "text"]
        )

        formatted_context = "\n\n".join(
            [f"{doc.entity.get('text')}" for doc in search_results[0]]
        )

        return str(formatted_context)

    except Exception as e:
        print(f"Error during the search operation: {e}")
        return "Sorry to find nothing"